#!/bin/bash
# Script to test different worker counts to find optimal configuration

BASE_CMD="PYTHONPATH= venv-np126/bin/python main.py \
    --train_data BSDS \
    --choose_test_data 1 \
    --input_dir /home/dd/Public/2440/ \
    --input_val_dir /home/dd/Public/2440/ \
    --output_dir checkpoints2440 \
    --train_list train_pair.lst \
    --test_list val_pair.lst \
    --disable_val_size_alignment \
    --test_img_height 2448 \
    --test_img_width 2080 \
    --crop_img true \
    --use_augmentation --aug_flip --aug_rotation --aug_brightness --aug_contrast --aug_noise \
    --epochs 1800 \
    --save_interval 5 \
    --batch_size 20 \
    --lr 1e-3 \
    --amp \
    --hy_data_dir /home/dd/Public/dexined_labeled/train/ \
    --hy_repeat 20 \
    --checkpoint_data /home/dd/Public/10_model_origin.pth"

echo "Testing with different worker counts..."
echo "Press Ctrl+C to stop if one works"
echo ""

# Test with recommended worker counts (from lowest to highest)
for WORKERS in 8 16 24 32 40; do
    echo "========================================="
    echo "Testing with --workers $WORKERS"
    echo "========================================="
    
    $BASE_CMD --workers $WORKERS
    
    EXIT_CODE=$?
    if [ $EXIT_CODE -eq 0 ] || [ $EXIT_CODE -eq 130 ]; then
        echo ""
        echo "SUCCESS with $WORKERS workers!"
        echo "Recommended: Use --workers $WORKERS"
        break
    else
        echo "Failed with $WORKERS workers (exit code: $EXIT_CODE)"
        echo "Trying next configuration..."
        sleep 2
    fi
done
